import os
import argparse
import random
from shutil import copy


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--dataset', default='optimal',
                        choices=['ucm', 'nwpu', 'aid', 'rs19', 'optimal'],
                        help='dataset name')
    parser.add_argument('--ratio', default=0.8, help="the ratio of train dataset")
    args = parser.parse_args()

    return args


def mkfile(file):
    if not os.path.exists(file):
        os.makedirs(file)


if __name__ == '__main__':
    args = parse_args()
    if args.dataset == 'ucm':
        path = "../dataset/ucm/Images"
        new_path = "../data/{0}".format("ucm" + "_" + str(args.ratio))
        mkfile(new_path + '/train')
        mkfile(new_path + '/val')
        categories = os.listdir(path)
        print(categories)
        for category in categories:
            mkfile(new_path + '/train/' + category)
            mkfile(new_path + '/val/' + category)
            category_path = os.path.join(path, category)
            images = os.listdir(category_path)
            num = len(images)

            train_index = random.sample(images, k=int(num * args.ratio))
            for index, image in enumerate(images):
                if image in train_index:
                    image_path = os.path.join(category_path, image)
                    new_path_ = new_path + '/train/' + category
                    copy(image_path, new_path_)
                else:
                    image_path = os.path.join(category_path, image)
                    new_path_ = new_path + '/val/' + category
                    copy(image_path, new_path_)
                print("\r[{}] processing [{}/{}]".format(category, index + 1, num), end="")  # processing bar

    if args.dataset == 'nwpu':
        path = "../dataset/NWPU-RESISC45"
        new_path = "../data/{0}".format("nwpu" + "_" + str(args.ratio))
        mkfile(new_path + '/train')
        mkfile(new_path + '/val')
        categories = os.listdir(path)
        print(categories)
        for category in categories:
            mkfile(new_path + '/train/' + category)
            mkfile(new_path + '/val/' + category)
            category_path = os.path.join(path, category)
            images = os.listdir(category_path)
            num = len(images)

            train_index = random.sample(images, k=int(num * args.ratio))
            for index, image in enumerate(images):
                if image in train_index:
                    image_path = os.path.join(category_path, image)
                    new_path_ = new_path + '/train/' + category
                    copy(image_path, new_path_)
                else:
                    image_path = os.path.join(category_path, image)
                    new_path_ = new_path + '/val/' + category
                    copy(image_path, new_path_)
                print("\r[{}] processing [{}/{}]".format(category, index + 1, num), end="")  # processi
    if args.dataset == 'aid':
        path = "../dataset/AID"
        new_path = "../data/{0}".format("aid" + "_" + str(args.ratio))
        mkfile(new_path + '/train')
        mkfile(new_path + '/val')
        categories = os.listdir(path)
        print(categories)
        for category in categories:
            mkfile(new_path + '/train/' + category)
            mkfile(new_path + '/val/' + category)
            category_path = os.path.join(path, category)
            images = os.listdir(category_path)
            num = len(images)

            train_index = random.sample(images, k=int(num * args.ratio))
            for index, image in enumerate(images):
                if image in train_index:
                    image_path = os.path.join(category_path, image)
                    new_path_ = new_path + '/train/' + category
                    copy(image_path, new_path_)
                else:
                    image_path = os.path.join(category_path, image)
                    new_path_ = new_path + '/val/' + category
                    copy(image_path, new_path_)
                print("\r[{}] processing [{}/{}]".format(category, index + 1, num), end="")  # processi

    if args.dataset == 'rs19':
        path = "../dataset/RS19"
        new_path = "../data/{0}".format("rs19" + "_" + str(args.ratio))
        mkfile(new_path + '/train')
        mkfile(new_path + '/val')
        categories = os.listdir(path)
        print(categories)
        for category in categories:
            mkfile(new_path + '/train/' + category)
            mkfile(new_path + '/val/' + category)
            category_path = os.path.join(path, category)
            images = os.listdir(category_path)
            num = len(images)

            train_index = random.sample(images, k=int(num * args.ratio))
            for index, image in enumerate(images):
                if image in train_index:
                    image_path = os.path.join(category_path, image)
                    new_path_ = new_path + '/train/' + category
                    copy(image_path, new_path_)
                else:
                    image_path = os.path.join(category_path, image)
                    new_path_ = new_path + '/val/' + category
                    copy(image_path, new_path_)
                print("\r[{}] processing [{}/{}]".format(category, index + 1, num), end="")  # processi
    if args.dataset == 'optimal':
        path = "../dataset/optimal"
        new_path = "../data/{0}".format("optimal" + "_" + str(args.ratio))
        mkfile(new_path + '/train')
        mkfile(new_path + '/val')
        categories = os.listdir(path)
        print(categories)
        for category in categories:
            mkfile(new_path + '/train/' + category)
            mkfile(new_path + '/val/' + category)
            category_path = os.path.join(path, category)
            images = os.listdir(category_path)
            num = len(images)

            train_index = random.sample(images, k=int(num * args.ratio))
            for index, image in enumerate(images):
                if image in train_index:
                    image_path = os.path.join(category_path, image)
                    new_path_ = new_path + '/train/' + category
                    copy(image_path, new_path_)
                else:
                    image_path = os.path.join(category_path, image)
                    new_path_ = new_path + '/val/' + category
                    copy(image_path, new_path_)
                print("\r[{}] processing [{}/{}]".format(category, index + 1, num), end="")  # processi
print("processing done!")
